from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from pydantic import BaseModel, Field
from typing import List


_gather_system_template = """
您是一个优秀的助手，根据参考信息,准确的完成用户查询目标的任务。
只回答与提问相关的答案，不要引入其它信息，精简内容不要过多的修饰。
"""
_gather_human_template = """
参考信息:
{infos}
查询目标:
{query}

"""


class Gather:
    def __init__(self, llm):
        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _gather_system_template),
                ("human", _gather_human_template),
            ]
        )

        _parser = StrOutputParser()
        self._chain = _prompt | llm | _parser

    def __call__(self, state):
        return self._chain.invoke(state)


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        base_url="http://192.168.10.11:60026/v1",
        model="qwen2.5:7b",
        api_key="ollama",
        temperature=0.4,
    )
    gather = Gather(_llm)
    _rt = gather(
        {
            "infos": ["全红婵是2024年法国巴黎奥运会女子10米跳台冠军", "她住在广东省湛江市"],
            "query": "2024年法国巴黎奥运会女子10米跳台冠军的家乡在哪里",
        }
    )
    print(_rt)
